import io as std_io
from io import BytesIO

import cv2
import numpy as np
import requests
from PIL import Image
from fastapi import FastAPI, Body
from fastapi.responses import JSONResponse
from modelscope.outputs import OutputKeys
from modules import script_callbacks

from scripts.kp.index_kp import kp_infer, cal_pos, merge
from scripts.lama.index_lama import *
from scripts.embedding.sam import *

log = logging.getLogger('sd')

inpainting = None
predictor = None  # 定义predictor为全局变量


# ai智能消除服务
def lama_inpaint_api(img: str = Body("a.jpg", title="orig img oss uri"),
                     mask: str = Body("b.jpg", title="img mask oss uri"), ):
    try:
        start_time = time.time()
        # 获取图片的url
        img_url = img
        mask_url = mask
        # 进行图像修复
        global inpainting  # 在这里也需要使用全局变量inpainting
        if inpainting is None:
            inpainting = load_pipeline()
        if inpainting is not None:
            result = inpainting({'img': img_url, 'mask': mask_url})
            log.info("inpainting finished : {:.6f} seconds".format(time.time() - start_time))  # 使用 logger 输出日志
            output_img = result[OutputKeys.OUTPUT_IMG]
            # 得到随机文件名
            file_name = get_png_name(img_url)
            # file_prefix = get_file_prefix()
            file_prefix = f"/mnt/output/{get_date_prefix()}/lama/"
            check_and_create_dir(file_prefix)
            file_path = os.path.join(file_prefix, file_name)

            response = requests.get(img_url)
            image_data = BytesIO(response.content)

            image = Image.open(image_data)
            alpha_channel = None
            if image.mode == "RGBA":
                np_img = np.array(image)
                alpha_channel = np_img[:, :, -1]

            if alpha_channel is not None:
                if alpha_channel.shape[:2] != output_img.shape[:2]:
                    alpha_channel = cv2.resize(alpha_channel, dsize=(output_img.shape[1], output_img.shape[0]))
                output_img = np.concatenate((output_img, alpha_channel[:, :, np.newaxis]), axis=-1)
            # 保存图片到本地
            cv2.imwrite(file_path, output_img)
            # 返回真实的oss地址
            log.info("inpainting finished all : {:.6f} seconds".format(time.time() - start_time))  # 使用 logger 输出日志
            result = f"hpk/{get_date_prefix()}/lama/{file_name}"
            return JSONResponse({'uri': result, "status": "success"})
    except Exception as e:
        return JSONResponse({'status': 'error', 'message': str(e)})


# infer人脸拼接服务
def infer_api(data: dict = Body(
    "https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/image_inpainting/image_inpainting.png",
    title="orig img oss url")):
    try:
        img_url = data.get("img")
        # 获取图片
        response = requests.get(img_url)
        #
        background = Image.open(std_io.BytesIO(response.content))
        # background = Image.open(image)
        head = Image.open('extensions/sd-extension-hpk/scripts/kp/head.png')
        #
        l, r = kp_infer(std_io.BytesIO(response.content))
        x, y, s = cal_pos(l, r)
        final_image = merge(background, head, x, y, s)
        # save image
        # date_str = datetime.datetime.now().strftime("%Y%m%d")
        output_dir = f"/mnt/output/{get_date_prefix()}/infer/"
        check_and_create_dir(output_dir)
        os.makedirs(output_dir, exist_ok=True)
        file_name = ''.join(random.choices(string.ascii_lowercase + string.digits, k=5)) + '.png'
        output_path = os.path.join(output_dir, file_name)
        final_image.save(output_path)
        result = f"hpk/{get_date_prefix()}/infer/{file_name}"
        return JSONResponse({'status': 'success', 'result': result})
    except Exception as e:
        return JSONResponse({'status': 'error', 'message': str(e)})


# 计算图片向量
def embedding_api(data: dict = Body(
    "https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/image_inpainting/image_inpainting.png",
    title="orig img oss url")):
    try:
        img_url = data['img']
        response = requests.get(img_url)
        #
        file_name = f"{get_random_string(4)}_{get_request_filename(img_url)}"
        image_path = f"{get_local_dir()}/{file_name}"
        # 保存图像到本地
        with open(image_path, 'wb') as f:
            f.write(response.content)
        # 加载图像
        image = cv2.imread(image_path)
        global predictor  # 在这里也需要使用全局变量predictor
        if predictor is None:
            print("Predictor is not available. Loading model...")
            try:
                predictor = load_predictor(model_type, checkpoint)
            except Exception as e:
                print("Failed to load model:", str(e))
                predictor = None
                return JSONResponse({'status': 'error', 'message': 'Failed to load model.'})
        if predictor is not None:
            predictor.set_image(image)
            # 将图像传递给预测器进行处理
            image_embedding = predictor.get_image_embedding().cuda().cpu().numpy()
            # image_path_prefix = get_file_prefix()
            image_path_prefix = f"/mnt/output/{get_date_prefix()}/embedding/"
            check_and_create_dir(image_path_prefix)
            embedding_file_name = f"{get_file_name_no_ext(file_name)}_embedding.npy"
            np.save(f"{image_path_prefix}/{embedding_file_name}", image_embedding)
            # predictor.reset_image()
            result = f"hpk/{get_date_prefix()}/embedding/{embedding_file_name}"
            return JSONResponse({'status': 'success', 'result': result})
        else:
            print("Failed to predict image.")
            return JSONResponse({'status': 'error', 'message': 'Failed to predict image.'})
    except Exception as e:
        return JSONResponse({'status': 'error', 'message': str(e)})
    finally:
        # 删除本地图像
        if os.path.exists(image_path):
            os.remove(image_path)


# api 注册
def register_api(app: FastAPI):
    app.add_api_route("/sdapi/v1/hpk/inpaint", lama_inpaint_api, methods=["POST"])
    app.add_api_route("/sdapi/v1/hpk/kf", infer_api, methods=["POST"])
    app.add_api_route("/sdapi/v1/hpk/embedding", embedding_api, methods=["POST"])


# 启动注册api
def on_app_started(blocks, app):  # register api
    register_api(app)


script_callbacks.on_app_started(on_app_started)
